package com.yc.config;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.yc.bean.Student;
import com.yc.services.ToolServices;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.output.structured.Description;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.*;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.elasticsearch.client.RestClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class SearchAiConfig {
    // 定义AI服务接口（注意保持命名一致）
    public interface SearchAiiAssistant {
//        @SystemMessage("你是一个帮学生查分数的AI，你要返回转化等级即可，加等级的解释，等级的解释就是缺勤次数，平时作业有四次，取平均数，课堂表现，期末测试，总评成绩，最后返回的格式是：等级：xxx，解释：xxx")
        @SystemMessage("你是一个帮学生查分数的AI,用户给你学号的姓名，你返回其Student对象")
        Student searchStudentdate(@UserMessage String question);
        @SystemMessage("给你传学生的数据，根据这个数据告诉他的等级和解释，最后返回的格式是：等级：xxx，解释：xxx")
        String answer(@UserMessage String question);
    }


//    @Bean
//    public EmbeddingStore<TextSegment> elasticsearchEmbeddingStore(ElasticsearchClient elasticsearchClient,
//                                                                   QwenEmbeddingModel qwenEmbeddingModel) {
//        // 从 ElasticsearchClient 中获取底层的 RestClient
//        RestClient restClient = ((RestClientTransport) elasticsearchClient._transport()).restClient();
//
//        return ElasticsearchEmbeddingStore.builder()
//                .restClient(restClient)
//                .indexName("langchain4j-embeddings")
//                .dimension(qwenEmbeddingModel.dimension())
//                .build();
//    }

    @Bean
    public SearchAiiAssistant searchAiiAssistant(ChatModel chatModel,
                                            StreamingChatModel streamingChatModel,
                                            ChatMemoryStore chatMemoryStore,
                                            ToolServices tools,
                                            EmbeddingStore embeddingStore,
                                            QwenEmbeddingModel qwenEmbeddingModel) {

        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(1000)
                //.chatMemoryStore(chatMemoryStore)
                .build();

        EmbeddingStoreContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(qwenEmbeddingModel)
                .embeddingStore(embeddingStore)
                .build();

        return AiServices.builder(SearchAiiAssistant.class)
                .chatModel(chatModel).streamingChatModel(streamingChatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .tools(tools)
                .contentRetriever(retriever)
                .build();
    }
}
